#!/usr/bin/env python3
"""
回测脚本

支持多交易所、多策略的标准化回测
"""

import sys
import os
import argparse
import json
import logging
from datetime import datetime, timedelta
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from xcquant.core import config_manager, unified_logger
from xcquant.data.data_manager import DataManager
from xcquant.strategies.enhanced_grid_strategy import EnhancedGridStrategy
from xcquant.exchanges.okx_exchange import OKXExchange
from xcquant.exchanges.binance_exchange import BinanceExchange
from xcquant.backtesting.backtest_engine import BacktestEngine
from xcquant.backtesting.mock_managers import MockOrderManager, MockAccountManager


def setup_logging(level=logging.INFO):
    """设置日志"""
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler('backtest.log')
        ]
    )
    return logging.getLogger(__name__)


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='量化交易回测工具')
    
    parser.add_argument('--config', '-c',
                       default='config.yaml',
                       help='配置文件路径')
    
    parser.add_argument('--strategy', '-s',
                       help='策略名称 (配置文件中定义)')
    
    parser.add_argument('--exchange', '-e',
                       choices=['okx', 'binance'],
                       help='交易所')
    
    parser.add_argument('--symbol',
                       help='交易对 (如: BTC-USDT)')
    
    parser.add_argument('--type', '-t',
                       choices=['spot', 'futures'],
                       help='交易类型')
    
    parser.add_argument('--start-date',
                       help='回测开始日期 (YYYY-MM-DD)')
    
    parser.add_argument('--end-date',
                       help='回测结束日期 (YYYY-MM-DD)')
    
    parser.add_argument('--initial-balance',
                       type=float,
                       help='初始资金')
    
    parser.add_argument('--output', '-o',
                       help='输出目录')
    
    parser.add_argument('--save-trades',
                       action='store_true',
                       help='保存交易记录')
    
    parser.add_argument('--save-charts',
                       action='store_true',
                       help='保存图表')
    
    parser.add_argument('--list-strategies',
                       action='store_true',
                       help='列出可用策略')
    
    parser.add_argument('--dry-run',
                       action='store_true',
                       help='试运行模式')
    
    return parser.parse_args()


def list_available_strategies(config: dict):
    """列出可用的策略"""
    strategies = config.get('strategies', {})
    
    print("\n可用策略:")
    for name, strategy_config in strategies.items():
        enabled = strategy_config.get('enabled', False)
        symbol = strategy_config.get('symbol', 'N/A')
        exchange = strategy_config.get('exchange', 'N/A')
        trade_type = strategy_config.get('type', 'N/A')
        status = "启用" if enabled else "禁用"
        
        print(f"  {name}:")
        print(f"    状态: {status}")
        print(f"    交易对: {symbol}")
        print(f"    交易所: {exchange}")
        print(f"    类型: {trade_type}")
        print()


def create_strategy(strategy_config: dict, exchange_adapter, 
                   order_manager, account_manager, data_manager):
    """创建策略实例"""
    strategy_type = strategy_config.get('strategy_type', 'grid')
    
    if strategy_type == 'grid':
        return EnhancedGridStrategy(
            strategy_config, exchange_adapter, 
            order_manager, account_manager, data_manager
        )
    else:
        raise ValueError(f"不支持的策略类型: {strategy_type}")


def create_exchange_adapter(exchange_name: str, config: dict):
    """创建交易所适配器"""
    exchange_config = config.get('exchanges', {}).get(exchange_name, {})
    
    if exchange_name == 'okx':
        return OKXExchange(exchange_config)
    elif exchange_name == 'binance':
        return BinanceExchange(exchange_config)
    else:
        raise ValueError(f"不支持的交易所: {exchange_name}")


def prepare_backtest_data(data_manager: DataManager, exchange: str, 
                         symbol: str, interval: str, trade_type: str,
                         start_date: datetime, end_date: datetime, logger):
    """准备回测数据"""
    logger.info(f"准备回测数据: {exchange} {symbol} {interval}")
    
    # 获取数据
    data = data_manager.get_data(
        exchange, symbol, interval, trade_type, 
        start_date, end_date
    )
    
    if data is None or data.empty:
        logger.error(f"无法获取回测数据: {exchange} {symbol}")
        return None
    
    logger.info(f"回测数据准备完成: {len(data)} 条记录")
    logger.info(f"数据时间范围: {data['datetime'].min()} - {data['datetime'].max()}")
    
    return data


def run_single_backtest(config: dict, strategy_name: str, 
                       strategy_config: dict, args, logger):
    """运行单个策略回测"""
    logger.info(f"开始回测策略: {strategy_name}")
    
    # 获取策略参数
    exchange = args.exchange or strategy_config.get('exchange')
    symbol = args.symbol or strategy_config.get('symbol')
    trade_type = args.type or strategy_config.get('type', 'spot')
    
    if not all([exchange, symbol]):
        logger.error("缺少必要参数: exchange, symbol")
        return None
    
    # 时间范围
    backtest_config = config.get('backtest', {})
    if args.start_date and args.end_date:
        start_date = datetime.strptime(args.start_date, '%Y-%m-%d')
        end_date = datetime.strptime(args.end_date, '%Y-%m-%d')
    else:
        end_date = datetime.now()
        days = backtest_config.get('default_days', 30)
        start_date = end_date - timedelta(days=days)
    
    # 初始资金
    initial_balance = args.initial_balance or backtest_config.get('initial_balance', 10000)
    
    logger.info(f"回测参数:")
    logger.info(f"  策略: {strategy_name}")
    logger.info(f"  交易所: {exchange}")
    logger.info(f"  交易对: {symbol}")
    logger.info(f"  类型: {trade_type}")
    logger.info(f"  时间范围: {start_date} - {end_date}")
    logger.info(f"  初始资金: {initial_balance}")
    
    if args.dry_run:
        logger.info("试运行模式，不执行实际回测")
        return None
    
    try:
        # 初始化数据管理器
        data_config = config.get('data', {})
        data_manager = DataManager(data_config)
        
        # 注册交易所
        exchange_adapter = create_exchange_adapter(exchange, config)
        data_manager.register_exchange(exchange, exchange_adapter)
        
        # 准备数据
        interval = strategy_config.get('interval', '1m')
        backtest_data = prepare_backtest_data(
            data_manager, exchange, symbol, interval, trade_type,
            start_date, end_date, logger
        )
        
        if backtest_data is None:
            return None
        
        # 创建回测引擎
        backtest_config = {
            'initial_balance': initial_balance,
            'commission_rate': 0.001,  # 手续费
            'slippage_rate': 0.0001,   # 滑点
            'start_date': start_date,
            'end_date': end_date
        }
        backtest_engine = BacktestEngine(backtest_config)
        
        # 创建模拟管理器
        order_manager = MockOrderManager()
        account_manager = MockAccountManager(initial_balance)
        
        strategy = create_strategy(
            strategy_config, exchange_adapter,
            order_manager, account_manager, data_manager
        )
        
        # 运行回测
        logger.info("开始执行回测...")
        results = backtest_engine.run_backtest(strategy, backtest_data)
        
        if not results:
            logger.error("回测执行失败")
            return None
        
        # 计算指标 (简化版本，因为BacktestMetrics可能不存在)
        performance = calculate_performance_metrics(results, initial_balance)
        
        # 输出结果
        logger.info("回测完成!")
        logger.info(f"总收益率: {performance['total_return']:.2%}")
        logger.info(f"最大回撤: {performance['max_drawdown']:.2%}")
        logger.info(f"交易次数: {performance['total_trades']}")
        logger.info(f"胜率: {performance['win_rate']:.2%}")
        
        # 保存结果
        output_dir = args.output or backtest_config.get('output_dir', './backtest_results')
        save_backtest_results(
            results, performance, strategy_name, 
            output_dir, args, logger
        )
        
        return {
            'strategy': strategy_name,
            'results': results,
            'performance': performance,
            'total_return': performance['total_return']
        }
        
    except Exception as e:
        logger.error(f"回测执行失败: {e}")
        return None


def calculate_performance_metrics(results: dict, initial_balance: float) -> dict:
    """计算回测性能指标"""
    trades = results.get('trades', [])
    portfolio_values = results.get('portfolio_values', [])
    
    if not portfolio_values:
        return {
            'total_return': 0.0,
            'max_drawdown': 0.0,
            'total_trades': len(trades),
            'win_rate': 0.0
        }
    
    # 计算总收益率
    final_value = portfolio_values[-1] if portfolio_values else initial_balance
    total_return = (final_value - initial_balance) / initial_balance
    
    # 计算最大回撤
    peak = initial_balance
    max_drawdown = 0.0
    
    for value in portfolio_values:
        if value > peak:
            peak = value
        drawdown = (peak - value) / peak
        if drawdown > max_drawdown:
            max_drawdown = drawdown
    
    # 计算胜率
    winning_trades = sum(1 for trade in trades if trade.get('profit', 0) > 0)
    win_rate = winning_trades / len(trades) if trades else 0.0
    
    return {
        'total_return': total_return,
        'max_drawdown': max_drawdown,
        'total_trades': len(trades),
        'win_rate': win_rate
    }


def save_backtest_results(results: dict, performance: dict, 
                         strategy_name: str, output_dir: str, 
                         args, logger):
    """保存回测结果"""
    try:
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 保存性能指标
        metrics_file = output_path / f"{strategy_name}_metrics_{timestamp}.json"
        with open(metrics_file, 'w', encoding='utf-8') as f:
            json.dump(performance, f, indent=2, ensure_ascii=False, default=str)
        
        logger.info(f"性能指标已保存: {metrics_file}")
        
        # 保存交易记录
        if args.save_trades and 'trades' in results:
            import pandas as pd
            trades_df = pd.DataFrame(results['trades'])
            trades_file = output_path / f"{strategy_name}_trades_{timestamp}.csv"
            trades_df.to_csv(trades_file, index=False)
            logger.info(f"交易记录已保存: {trades_file}")
        
        # 保存图表
        if args.save_charts:
            save_backtest_charts(results, performance, strategy_name, 
                               output_path, timestamp, logger)
        
    except Exception as e:
        logger.error(f"保存回测结果失败: {e}")


def save_backtest_charts(results: dict, performance: dict, 
                        strategy_name: str, output_path: Path, 
                        timestamp: str, logger):
    """保存回测图表"""
    try:
        import matplotlib.pyplot as plt
        import pandas as pd
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle(f'{strategy_name} 回测结果', fontsize=16)
        
        # 资金曲线
        if 'equity_curve' in results:
            equity_df = pd.DataFrame(results['equity_curve'])
            axes[0, 0].plot(equity_df['datetime'], equity_df['equity'])
            axes[0, 0].set_title('资金曲线')
            axes[0, 0].set_xlabel('时间')
            axes[0, 0].set_ylabel('资金')
            
        # 回撤曲线
        if 'drawdown' in results:
            drawdown_df = pd.DataFrame(results['drawdown'])
            axes[0, 1].fill_between(drawdown_df['datetime'], 0, drawdown_df['drawdown'], 
                                   alpha=0.3, color='red')
            axes[0, 1].set_title('回撤曲线')
            axes[0, 1].set_xlabel('时间')
            axes[0, 1].set_ylabel('回撤 (%)')
        
        # 收益分布
        if 'trades' in results:
            trades_df = pd.DataFrame(results['trades'])
            if not trades_df.empty and 'pnl' in trades_df.columns:
                axes[1, 0].hist(trades_df['pnl'], bins=30, alpha=0.7)
                axes[1, 0].set_title('收益分布')
                axes[1, 0].set_xlabel('单笔收益')
                axes[1, 0].set_ylabel('频次')
        
        # 关键指标
        metrics_text = f"""
        总收益率: {performance.get('total_return', 0):.2%}
        年化收益率: {performance.get('annual_return', 0):.2%}
        最大回撤: {performance.get('max_drawdown', 0):.2%}
        夏普比率: {performance.get('sharpe_ratio', 0):.4f}
        交易次数: {performance.get('total_trades', 0)}
        胜率: {performance.get('win_rate', 0):.2%}
        """
        
        axes[1, 1].text(0.1, 0.5, metrics_text, fontsize=12, 
                        verticalalignment='center')
        axes[1, 1].set_title('关键指标')
        axes[1, 1].axis('off')
        
        plt.tight_layout()
        
        # 保存图表
        chart_file = output_path / f"{strategy_name}_chart_{timestamp}.png"
        plt.savefig(chart_file, dpi=300, bbox_inches='tight')
        plt.close()
        
        logger.info(f"回测图表已保存: {chart_file}")
        
    except Exception as e:
        logger.error(f"保存回测图表失败: {e}")


def main():
    """主函数"""
    args = parse_arguments()
    
    # 设置日志
    logger = setup_logging()
    
    # 加载配置
    try:
        config_manager = ConfigManager(args.config)
        config = config_manager.get_config()
        logger.info(f"配置文件加载成功: {args.config}")
    except Exception as e:
        logger.error(f"加载配置文件失败: {e}")
        sys.exit(1)
    
    # 列出策略
    if args.list_strategies:
        list_available_strategies(config)
        return
    
    # 获取策略配置
    strategies = config.get('strategies', {})
    
    if args.strategy:
        # 运行指定策略
        if args.strategy not in strategies:
            logger.error(f"策略不存在: {args.strategy}")
            sys.exit(1)
        
        strategy_config = strategies[args.strategy].copy()
        
        # 命令行参数覆盖配置
        if args.exchange:
            strategy_config['exchange'] = args.exchange
        if args.symbol:
            strategy_config['symbol'] = args.symbol
        if args.type:
            strategy_config['type'] = args.type
        
        try:
            result = run_single_backtest(config, args.strategy, strategy_config, args, logger)
            
            if result:
                logger.info("回测完成")
            else:
                logger.error("回测失败")
                sys.exit(1)
        except Exception as e:
            logger.error(f"回测执行失败: {e}")
            sys.exit(1)
    
    else:
        # 运行所有启用的策略
        enabled_strategies = {
            name: cfg for name, cfg in strategies.items() 
            if cfg.get('enabled', False)
        }
        
        if not enabled_strategies:
            logger.error("没有启用的策略")
            sys.exit(1)
        
        logger.info(f"将运行 {len(enabled_strategies)} 个策略")
        
        results = []
        failed_strategies = []
        
        for strategy_name, strategy_config in enabled_strategies.items():
            try:
                result = run_single_backtest(config, strategy_name, strategy_config, args, logger)
                if result:
                    results.append(result)
                else:
                    failed_strategies.append(strategy_name)
            except Exception as e:
                logger.error(f"策略 {strategy_name} 执行失败: {e}")
                failed_strategies.append(strategy_name)
        
        # 输出汇总
        if results:
            logger.info(f"成功完成 {len(results)} 个策略回测")
            
            # 汇总统计
            total_return = sum(r.get('total_return', 0) for r in results)
            avg_return = total_return / len(results)
            
            logger.info(f"平均收益率: {avg_return:.2%}")
            
        if failed_strategies:
            logger.warning(f"失败的策略: {', '.join(failed_strategies)}")
            
        if not results:
            logger.error("所有策略回测都失败了")
            sys.exit(1)
            logger.info(f"\n回测汇总 ({len(results)} 个策略):")
            for result in results:
                perf = result['performance']
                logger.info(f"  {result['strategy']}: "
                          f"收益率 {perf.get('total_return', 0):.2%}, "
                          f"最大回撤 {perf.get('max_drawdown', 0):.2%}")
        else:
            logger.error("所有策略回测都失败了")
            sys.exit(1)


if __name__ == '__main__':
    main()